import numpy as np
import torchvision
import time
import os
import copy
import pdb
import time
import argparse

import sys
import cv2

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, models, transforms
from retinanet import model
from retinanet.dataloader import  CSVDataset, collater, Resizer, AspectRatioBasedSampler, Augmenter, \
    UnNormalizer, Normalizer

assert torch.__version__.split('.')[0] == '1'

print('CUDA available: {}'.format(torch.cuda.is_available()))


iou_thresh = 0.3

score_thresh = 0.1
Save_Path = '../../data/show/baseline_person'
os.makedirs(Save_Path,exist_ok=True)

def main(args=None):
    parser = argparse.ArgumentParser(description='Simple training script for training a RetinaNet network.')
    parser.add_argument('--dataset', help='Dataset type, must be one of csv or coco.', default='csv')
    parser.add_argument('--coco_path', help='Path to COCO directory')
    parser.add_argument('--csv_val', help='Path to COCO directory', default='../label_file/val_resized_person.csv')
    parser.add_argument('--model_path', help='Path to model', type=str,default='../../model_save/baseline_person/')
    parser.add_argument('--csv_classes', help='Path to file containing class list (see readme)', default='../label_file/class_person.csv')
    parser.add_argument('--depth', help='Resnet depth, must be one of 18, 34, 50, 101, 152', default='new_efficientdet_b3')

    parser = parser.parse_args(args)

    dataset_val = CSVDataset(train_file=parser.csv_val, class_list=parser.csv_classes,
                             transform=transforms.Compose([Normalizer(), Resizer()]))

    sampler_val = AspectRatioBasedSampler(dataset_val, batch_size=1, drop_last=False)
    dataloader_val = DataLoader(dataset_val, num_workers=1, collate_fn=collater, batch_sampler=None)

    # Create the model
    if parser.depth == 18:
        retinanet = model.resnet18(num_classes=dataset_val.num_classes(), pretrained=True)
    elif parser.depth == 34:
        retinanet = model.resnet34(num_classes=dataset_val.num_classes(), pretrained=True)
    elif parser.depth == 50:
        retinanet = model.resnet50(num_classes=dataset_val.num_classes(), pretrained=True)
    elif parser.depth == 101:
        retinanet = model.resnet101(num_classes=dataset_val.num_classes(), pretrained=True)
    elif parser.depth == 152:
        retinanet = model.resnet152(num_classes=dataset_val.num_classes(), pretrained=True)
    elif parser.depth == 'efficientnet-b0':
        retinanet = model.resnet18(pretrained=True, num_classes=dataset_val.num_classes(),
                                   effici_name='efficientnet-b0')
    elif parser.depth == 'new_efficientdet_b3':
        retinanet = model.new_efficientdet_b3(pretrained=True, num_classes=dataset_val.num_classes(),thresh=iou_thresh)
    else:
        raise ValueError('Unsupported model depth, must be one of 18, 34, 50, 101, 152')

    retinanet = torch.nn.DataParallel(retinanet).cuda()
    # load trained model weights
    trained_model_state_dict = torch.load(parser.model_path)
    
    # trained_model_state_dict = {k.replace('module.', ''): v for k, v in trained_model_state_dict.items()}
    retinanet.load_state_dict(trained_model_state_dict)

    use_gpu = torch.cuda.is_available()
    if use_gpu:
        retinanet = retinanet.cuda()

    retinanet.eval()

    unnormalize = UnNormalizer()

    def draw_caption_label(image, box, caption):

        b = np.array(box).astype(int)
        ## 用黑色和白色分别写一次，白色较细在黑色上面，容易显示
        cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2)
        cv2.putText(image, caption, (b[0], b[1] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)

    def draw_caption_pred(image, box, caption):

        b = np.array(box).astype(int)
        ## 用黑色和白色分别写一次，白色较细在黑色上面，容易显示
        cv2.putText(image, caption, (b[2], b[3] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (0, 0, 0), 2)
        cv2.putText(image, caption, (b[2], b[3] - 10), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), 1)

    for idx, data in enumerate(dataloader_val):

        with torch.no_grad():

            scores, classification, transformed_anchors = retinanet(data['img'].cuda().float())
            scores = scores[:6 ]
            label_annot = (data['annot'])[0]
            idxs = np.where(scores.cpu() > score_thresh)
            img = np.array(255 * unnormalize(data['img'][0, :, :, :])).copy()

            img[img < 0] = 0
            img[img > 255] = 255

            img = np.transpose(img, (1, 2, 0))

            img = cv2.cvtColor(img.astype(np.uint8), cv2.COLOR_BGR2RGB)

            # 逐个画出预测到置信度符合阈值的bbox

            for j in range(idxs[0].shape[0]):
                bbox = transformed_anchors[idxs[0][j], :]
                x1 = int(bbox[0])
                y1 = int(bbox[1])
                x2 = int(bbox[2])
                y2 = int(bbox[3])


                draw_caption_pred(img, (x1, y1, x2, y2), ' ')
                cv2.rectangle(img, (x1, y1), (x2, y2), color=(0, 0, 255), thickness=1)

            ### 逐个画出label上的bbox
            for i in range(label_annot.shape[0]):
                label_x1 = int(label_annot[i][0])
                label_y1 = int(label_annot[i][1])
                label_x2 = int(label_annot[i][2])
                label_y2 = int(label_annot[i][3])


                draw_caption_label(img, (label_x1, label_y1, label_x2, label_y2), ' ')
                cv2.rectangle(img, (label_x1, label_y1), (label_x2, label_y2), color=(0, 255, 255), thickness=1)
            
        cv2.imwrite((Save_Path+'/'+str(idx)+'.png'),img)
        cv2.imshow('img', img)
        cv2.waitKey(0)


if __name__ == '__main__':
    main()